package org.pine.start.interceptor;

import cn.hutool.cache.impl.LRUCache;
import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.core.metadata.TableFieldInfo;
import com.baomidou.mybatisplus.core.metadata.TableInfo;
import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.expression.Alias;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.HexValue;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.select.PlainSelect;
import org.pine.api.enums.ResultCodeEnum;
import org.pine.common.annotation.DataPermission;
import org.pine.common.annotation.DataPermissionGroup;
import org.pine.common.beans.UserInfo;
import org.pine.common.enums.DbDataScopeEnum;
import org.pine.common.exception.PineException;
import org.pine.common.utils.UserUtil;

import java.lang.reflect.Method;
import java.util.*;

/**
 * 角色数据权限处理器
 *
 * @author pine manage
 * @since 2024-09-26
 */
@Slf4j
public class RoleCustomDataPermissionHandler implements CustomDataPermissionHandler {

    private static final Map<DbDataScopeEnum, ProcessWhereExpressionListener> PROCESS_WHERE_EXPRESSION_LISTENER_MAP = new HashMap<>(8);

    /**
     * 表名-字段名缓存
     */
    private static final Map<String, Set<String>> TABLE_FIELD_NAME_CACHE = new HashMap<>(16);

    /**
     * 数据权限组缓存：通过解析mapper接口内标注了DataPermissionGroup或DataPermission注解的方法获取数据权限组配置
     */
    private static final LRUCache<String, List<DataPermission>> DATA_PERMISSION_GROUP_CACHE = new LRUCache<>(100);

    @Override
    @SneakyThrows(Exception.class)
    public Expression getSqlSegmentWithPermission(PlainSelect plainSelect, String whereStatement) {
        log.info("开始进行权限过滤{}", whereStatement);
        // 待执行 SQL Where 条件表达式
        Expression where = plainSelect.getWhere();
        UserInfo userInfo = UserUtil.get();
        if (Objects.isNull(userInfo)) {
            log.info("当前无登录用户信息，不处理数据权限");
            return where;
        }

        if (where == null) {
            where = new HexValue(" 1 = 1 ");
        }
        // 获取mapper名称
        String className = whereStatement.substring(0, whereStatement.lastIndexOf("."));
        // 获取方法名
        String methodName = whereStatement.substring(whereStatement.lastIndexOf(".") + 1);
        Table fromItem = (Table) plainSelect.getFromItem();
        // 有别名用别名，无别名用表名，防止字段冲突报错
        Alias fromItemAlias = fromItem.getAlias();
        String mainTableName = fromItemAlias == null ? fromItem.getName() : fromItemAlias.getName();
        // 获取当前mapper 的方法
        Method[] methods = Class.forName(className).getMethods();
        // 刷新表字段缓存
        refreshTableFieldNameCache(fromItem.getName());

        // 遍历mapper
        for (Method method : methods) {
            if (Objects.equals(method.getName(), methodName)) {
                boolean annotationGroupPresent = method.isAnnotationPresent(DataPermissionGroup.class);
                boolean annotationPresent = method.isAnnotationPresent(DataPermission.class);
                if (!(annotationPresent || annotationGroupPresent)) {
                    log.info("当前方法无权限注解，跳过权限处理");
                    return where;
                }

                log.info("权限处理开始，执行方法：{}", whereStatement);
                try {
                    // 获取当前角色的数据权限信息
                    DbDataScopeEnum roleDataScope = getRoleDataScope();
                    if (Objects.isNull(roleDataScope)) {
                        roleDataScope = DbDataScopeEnum.ONLY_SELF;
                    }

                    // 解析注解
                    DataPermission dataPermission = parseAndGetDataPermission(whereStatement, method, roleDataScope);
                    if (Objects.isNull(dataPermission)) {
                        // 权限注解中，没有匹配数据类型的权限信息
                        log.info("当前方法未配置权限可生效的DataPermission注解，跳过权限处理");
                        return where;
                    }

                    if (DbDataScopeEnum.ALL.equals(dataPermission.dataScope())) {
                        // 权限注解中，没有匹配数据类型的权限信息
                        log.info("当前方法配置权限为所有数据，跳过权限处理");
                        return where;
                    }

                    String permissionFieldName = dataPermission.permissionFieldName();
                    if (StrUtil.isBlank(permissionFieldName)) {
                        throw new PineException(ResultCodeEnum.ARGS_MISSING_EXCEPTION.getCode(), "权限字段名不能为空");
                    }
                    Set<String> tableFieldNames = TABLE_FIELD_NAME_CACHE.get(fromItem.getName());
                    if (!tableFieldNames.contains(permissionFieldName)) {
                        log.info("权限字段匹配，表字段名未命中，跳过权限处理，表名: {}, 字段名: {}", fromItem.getName(), permissionFieldName);
                        return where;
                    }

                    log.info("权限字段匹配，表字段名缓存命中，权限作用域：{}, 表名: {}, 字段名: {}", roleDataScope, fromItem.getName(), permissionFieldName);
                    return processWhere(roleDataScope, where, dataPermission, mainTableName);
                } catch (Exception e) {
                    log.error("权限处理失败", e);
                    return new HexValue(" 1 = 2 ");
                }
            }
        }

        // 说明无权查看，
        where = new HexValue(" 1 = 2 ");
        return where;
    }

    /**
     * 获取当前角色的数据权限信息
     */
    protected DbDataScopeEnum getRoleDataScope() {
        UserInfo userInfo = UserUtil.get();
        return userInfo.getDataScope();
    }

    protected Expression processWhere(DbDataScopeEnum roleDataScope, Expression where, DataPermission dataPermission, String mainTableName) {
        ProcessWhereExpressionListener expressionListener = PROCESS_WHERE_EXPRESSION_LISTENER_MAP.get(roleDataScope);
        if (Objects.isNull(expressionListener)) {
            throw new PineException(ResultCodeEnum.DATA_EXCEPTION.getCode(), "未找到对应的数据权限处理器");
        }

        return expressionListener.process(where, dataPermission, mainTableName);
    }

    public static void registerProcessWhereExpressionListener(DbDataScopeEnum roleDataScope, ProcessWhereExpressionListener expressionListener) {
        log.info("注册权限处理器,roleDataScope: {},expressionListener: {}", roleDataScope, expressionListener);
        PROCESS_WHERE_EXPRESSION_LISTENER_MAP.put(roleDataScope, expressionListener);
    }

    private DataPermission parseAndGetDataPermission(String statementId, Method method, DbDataScopeEnum roleDataScope) {
        // 从缓存中获取注解权限配置
        if (DATA_PERMISSION_GROUP_CACHE.containsKey(statementId)) {
            return DATA_PERMISSION_GROUP_CACHE.get(statementId).stream()
                    .filter(dataPermission -> Objects.equals(dataPermission.dataScope(), roleDataScope))
                    .findFirst()
                    .orElse(null);
        }

        // 处理注解：依据当前的数据权限类型，获取对应的权限配置
        DataPermissionGroup dataPermissionGroup = method.getAnnotation(DataPermissionGroup.class);
        // 权限组为空
        if (Objects.isNull(dataPermissionGroup)) {
            DataPermission dataPermission = method.getAnnotation(DataPermission.class);
            DATA_PERMISSION_GROUP_CACHE.put(statementId, Collections.singletonList(dataPermission));
            return Objects.equals(dataPermission.dataScope(), roleDataScope) ? dataPermission : null;
        }

        // 处理权限组注解配置信息
        DataPermission[] dataPermissions = dataPermissionGroup.value();
        DATA_PERMISSION_GROUP_CACHE.put(statementId, Arrays.asList(dataPermissions));
        for (DataPermission currentDataPermission : dataPermissions) {
            if (Objects.equals(currentDataPermission.dataScope(), roleDataScope)) {
                return currentDataPermission;
            }
        }
        return null;
    }

    /**
     * 刷新表字段名缓存
     *
     * @param tableName 表名
     */
    private void refreshTableFieldNameCache(String tableName) {
        if (TABLE_FIELD_NAME_CACHE.containsKey(tableName)) {
            return;
        }

        // 获取表信息
        TableInfo tableInfo = TableInfoHelper.getTableInfo(tableName);
        // 获取表字段信息
        List<TableFieldInfo> fieldList = tableInfo.getFieldList();
        Set<String> fieldNameSet = new HashSet<>(fieldList.size());
        for (TableFieldInfo tableFieldInfo : fieldList) {
            // 字段名
            String column = tableFieldInfo.getColumn();
            fieldNameSet.add(column);
        }
        TABLE_FIELD_NAME_CACHE.put(tableName, fieldNameSet);
    }
}
